import os
from openai import OpenAI

# -------------------- 基础配置 --------------------
# API Key 建议放到系统环境变量中：DASHSCOPE_API_KEY
# 或者在这里直接写入，例如：
# api_key = "sk-xxxxxxxxxxxx"
client = OpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# Qiskit源目录与MindQuantum目标目录
SRC_DIR = r"E:\python-projects\qiskit2mq\data\qiskit_tasks"
DST_DIR = r"E:\python-projects\qiskit2mq\data\mq_tasks"

# -------------------- 翻译提示模板 --------------------
PROMPT_TEMPLATE = """You are an expert in quantum programming.

Translate the following Qiskit Python code into **MindQuantum** code.
- Preserve all function names, arguments, and structure.
- Replace Qiskit classes and functions with MindQuantum equivalents (Circuit, Simulator, etc.).
- Keep the logic and expected behavior consistent.
- Do not add extra explanations, only output valid Python code.

Qiskit code:
{code}
"""

# -------------------- 调用 Qwen3-Max 翻译 --------------------
def translate_code(file_path):
    """调用 qwen3-max 翻译单个文件"""
    with open(file_path, "r", encoding="utf-8") as f:
        code = f.read()

    prompt = PROMPT_TEMPLATE.format(code=code)

    completion = client.chat.completions.create(
        model="qwen3-max",
        messages=[
            {"role": "system", "content": "You are a professional quantum software translator."},
            {"role": "user", "content": prompt},
        ],
        temperature=0.3,
    )

    # 获取生成文本
    translated_code = completion.choices[0].message.content

    # 去掉可能的 Markdown 包裹
    if translated_code.startswith("```"):
        translated_code = translated_code.strip("`")
        translated_code = translated_code.replace("python\n", "").replace("Python\n", "")

    return translated_code


# -------------------- 主流程 --------------------
def main():
    py_files = sorted([f for f in os.listdir(SRC_DIR) if f.endswith(".py")])
    for filename in py_files:
        src_path = os.path.join(SRC_DIR, filename)
        dst_path = os.path.join(DST_DIR, filename)

        print(f"🔹 Translating {filename} ...")

        try:
            mq_code = translate_code(src_path)

            with open(dst_path, "w", encoding="utf-8") as f:
                f.write(mq_code)

            print(f"✅ Saved MindQuantum code to {dst_path}\n")
        except Exception as e:
            print(f"❌ Failed to translate {filename}: {e}\n")


if __name__ == "__main__":
    main()
